from flask import Flask, request, jsonify
from flask_cors import CORS
import jupyter_client


def create_app() -> Flask:
    app = Flask(__name__)
    CORS(app)

    # Start a single Python kernel the API can use
    km = jupyter_client.KernelManager()
    km.start_kernel()
    kc = km.client()
    kc.start_channels()

    @app.route('/health', methods=['GET'])
    def health() -> tuple:
        return jsonify({ 'status': 'ok' }), 200

    @app.route('/run', methods=['POST'])
    def run():
        payload = request.get_json(silent=True) or {}
        code = payload.get('code', '')
        if not code:
            return jsonify({ 'error': 'Missing code' }), 400

        msg_id = kc.execute(code)
        output_lines = []
        # Collect outputs until the kernel goes idle for this execution
        while True:
            msg = kc.get_iopub_msg()
            if msg.get('parent_header', {}).get('msg_id') != msg_id:
                continue
            msg_type = msg.get('msg_type')
            content = msg.get('content', {})

            if msg_type == 'stream':
                output_lines.append(content.get('text', ''))
            elif msg_type == 'execute_result':
                data = content.get('data', {})
                output_lines.append(data.get('text/plain', ''))
            elif msg_type == 'display_data':
                data = content.get('data', {})
                output_lines.append(data.get('text/plain', ''))
            elif msg_type == 'error':
                output_lines.append('\n'.join(content.get('traceback', [])))
            elif msg_type == 'status' and content.get('execution_state') == 'idle':
                break

        return jsonify({ 'result': ''.join(output_lines) })

    return app


if __name__ == '__main__':
    app = create_app()
    app.run(host='0.0.0.0', port=8888, debug=True)


